import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import h5py
import skimage.transform as tf
sys.path.append(os.path.abspath(os.path.dirname(__file__)) + "/..")


def load_dataset():
    '''
    加载数据集
    '''
    train_dataset = h5py.File('datasets/train_catvnoncat.h5', 'r')
    train_set_x = np.array(train_dataset['train_set_x'][:])
    train_set_y = np.array(train_dataset['train_set_y'][:])
    test_dataset = h5py.File('datasets/test_catvnoncat.h5', 'r')
    test_set_x = np.array(test_dataset['test_set_x'][:])
    test_set_y = np.array(test_dataset['test_set_y'][:])
    return train_set_x, train_set_y, test_set_x, test_set_y


train_set_x, train_set_y, test_set_x, test_set_y = load_dataset()
m_train = train_set_x.shape[0]
m_test = test_set_x.shape[0]
numpx = train_set_x.shape[1]
print(f'train_set_x {train_set_x.shape} train_set_y {train_set_y.shape}, test_set_x {test_set_x.shape} test_set_y {test_set_y.shape}')


def init_wb_parameters(n_x, n_h=4, n_y=1):
    '''
    初始化权重和偏置
    n_x:输入层个数
    n_h:隐藏层个数
    n_y:输出层个数
    注:不能为全0,用随机值
    '''
    np.random.seed(2)
    W1 = np.random.randn(n_h, n_x)*0.01
    B1 = np.zeros((n_h, 1))
    W2 = np.random.randn(n_y, n_h)*0.01
    B2 = np.zeros((n_y, 1))
    return {
        'W1': W1,
        'B1': B1,
        'W2': W2,
        'B2': B2,
    }


def sigmoid(Z):
    '''
    激活函数
    '''
    return 1/(1 + np.exp(-Z))


def propagate(parameters, X, Y):
    '''
    传播函数
    '''
    W1 = parameters['W1']
    B1 = parameters['B1']
    W2 = parameters['W2']
    B2 = parameters['B2']
    # 向前传播得到成本值
    # 隐藏层使用tanh激活函数
    A1 = np.tanh(np.dot(W1, X)+B1)
    # 输出层使用sigmoid激活函数
    A2 = sigmoid(np.dot(W2, A1) + B2)
    # 成本计算
    m = Y.shape[1]
    cost = -np.sum((Y*np.log(A2) + (1 - Y)*np.log(1-A2)))/m
    # 向后传播得到dw，db
    DZ2 = A2 - Y
    DW2 = np.dot(A1, DZ2.T)/m
    DB2 = np.sum(DZ2, axis=1, keepdims=True)/m
    DZ1 = np.multiply(np.dot(W2.T, DZ2), 1 - np.power(A1, 2))
    DW1 = np.dot(X, DZ1.T)/m
    DB1 = np.sum(DZ1, axis=1, keepdims=True)/m
    # print(f'DZ2 {DZ2.shape} DW2 {DW2.shape} DB2 {DB2.shape} DW1 {DW1.shape} DB1 {DB1.shape}')
    return {
        'DW2': DW2.T,
        'DB2': DB2,
        'DW1': DW1.T,
        'DB1': DB1,
    }, cost


def update_parameters(parameters, dx_parameters, learning_rate):
    '''
    更新dw参数
    '''
    W1 = parameters['W1']
    B1 = parameters['B1']
    W2 = parameters['W2']
    B2 = parameters['B2']
    DW1 = dx_parameters['DW1']
    DB1 = dx_parameters['DB1']
    DW2 = dx_parameters['DW2']
    DB2 = dx_parameters['DB2']
    # print(f'learning_rate {learning_rate} W1 {W1.shape} DW1 {DW1.shape}')
    W1 -= learning_rate*DW1
    B1 -= learning_rate*DB1
    W2 -= learning_rate*DW2
    B2 -= learning_rate*DB2
    return {
        'W1': W1,
        'B1': B1,
        'W2': W2,
        'B2': B2,
    }


def optimize(parameters, train_X, train_Y, num, learning_rate):
    '''
    优化函数，只需要w,b原始值及训练特征，训练标签，次数和学习率
    '''
    for i in range(num):
        # 开始传播
        dx_parameters, cost = propagate(parameters, train_X, train_Y)
        parameters = update_parameters(
            parameters, dx_parameters, learning_rate)
        if i % 100 == 0:
            print(f'第{i}次后成本 是 {cost}')
    return parameters


def predict(parameters, X):
    '''
    对特征进行预测
    '''
    prediction = np.zeros((1, X.shape[1]))
    W1 = parameters['W1']
    B1 = parameters['B1']
    W2 = parameters['W2']
    B2 = parameters['B2']
    # 向前传播得到成本值
    # 隐藏层使用tanh激活函数
    A1 = np.tanh(np.dot(W1, X)+B1)
    # 输出层使用sigmoid激活函数
    A2 = sigmoid(np.dot(W2, A1) + B2)
    for i in range(A2.shape[1]):
        t = A2[0, i]
        if t >= 0.5:
            prediction[0, i] = 1
    return prediction


def model(train_X, train_Y, test_X, test_Y, num=5000, learning_rate=0.01):
    '''
    模型训练
    '''
    #
    # 先得到wb
    parameters = init_wb_parameters(train_X.shape[0])
    # 优化得到理好的w,b
    parameters = optimize(parameters, train_X, train_Y, num, learning_rate)
    # 得到训练及测试预测
    train_prediction = predict(parameters, train_X)
    test_prediction = predict(parameters, test_X)
    # 打印出预测的准确率
    print(
        f'对训练图片的预测准确率为 {100 - np.mean(np.abs(train_prediction - train_Y))*100}')
    print(
        f'对测试图片的预测准确率为 {100 - np.mean(np.abs(test_prediction - test_Y))*100}')
    return parameters


# 转换为0-1间
train_set_x_flatten = train_set_x/255
test_set_x_flatten = test_set_x/255
# 再转换为易计算的n*m矩阵
train_set_X = train_set_x_flatten.reshape(train_set_x_flatten.shape[0], -1).T
train_set_Y = train_set_y.reshape(1, train_set_y.shape[0])
test_set_X = test_set_x_flatten.reshape(test_set_x_flatten.shape[0], -1).T
test_set_Y = test_set_y.reshape(1, test_set_y.shape[0])
parameters = model(train_set_X, train_set_Y, test_set_X, test_set_Y)
my_image = "my_image3.jpg"
fname = "images/" + my_image
image = np.array(plt.imread(fname))
my_image = tf.resize(image, (numpx, numpx),
                     mode='reflect').reshape((1, numpx*numpx*3)).T
my_predicted_image = predict(parameters, my_image)
print("预测结果为 " + str(int(np.squeeze(my_predicted_image))))
plt.imshow(image)
plt.show()
